import torch
import torch.optim as optim
import pandas as pd
from model import RecommenderModel
import torch.nn as nn
import config


class train_model():
    # Load the training data
    train_data = pd.read_csv('train.csv')

    # Extract inputs and targets from the train_data
    inputs = train_data[['userId', 'movieId']].values
    targets = train_data['rating'].values

    # Initialize the model
    model = RecommenderModel(config.model_input_size, config.model_hidden_size, config.model_output_size)

    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    for epoch in range(10):
        optimizer.zero_grad()
        outputs = model(torch.Tensor(inputs))
        loss = criterion(outputs.squeeze(), torch.Tensor(targets))  # Assuming the model output is a single value
        loss.backward()
        optimizer.step()

    # Save the trained model
    torch.save(model.state_dict(), 'model.pth')
